from Neural_Weight_Quantization.dataset_partition_1 import kmeans_multiple
from Neural_Weight_Quantization.util import dir_io
import time
import numpy as np
import copy
import os


def dataset_partition(base, base_base_gnd, config):
    dir_io.mkdir(config['program_train_para_dir'])
    start_time = time.time()
    partition_method = factory(config['dataset_partition']['type'])
    for i in range(config['n_classifier']):
        save_dir = '%s/Classifier_%d/' % (config['program_train_para_dir'], i)
        os.mkdir(save_dir)
        save_dir = '%s/Classifier_%d/dataset_partition' % (config['program_train_para_dir'], i)
        os.mkdir(save_dir)
    print("start dataset partition")
    partition_info_l, intermediate = partition_method(base, base_base_gnd, config)
    print("end dataset partition")
    end_time = time.time()

    intermediate['total_dataset_partition_time'] = end_time - start_time
    return partition_info_l, intermediate


def factory(_type):
    if _type == 'kmeans_multiple':
        return kmeans_multiple.dataset_partition
    raise Exception('do not support the type of partition')
